// pages/ai-assistant/index.js
import { createStoreBindings } from 'mobx-miniprogram-bindings'
import { tokenStore } from '../../stores/tokenStore'
const app = getApp()

Page({
  data: {
    // 聊天消息列表
    messages: [],
    // 输入框内容
    inputValue: "",
    // 是否正在加载回复
    loading: false,
    // 当前正在生成的回复内容
    currentResponse: "",
    // 当前正在生成的解析后内容
    currentParsedResponse: {},
    // WebSocket连接状态
    socketConnected: false,
    // 重连次数
    reconnectCount: 0,
    // 最大重连次数
    maxReconnectCount: 5,
    // 是否正在停止生成
    isStopping: false,
    // 滚动位置
    scrollTop: 0,
  },

  onLoad() {
    // 注入token信息
    this.tokenStoreBinding = createStoreBindings(this, {
      store: tokenStore,
      fields: ['token'],
      actions: ['updateToken']
    })
    // 连接WebSocket
    this.connectSocket()
  },

  onUnload() {
    // 页面卸载时关闭WebSocket连接
    this.closeSocket()
  },

  // 连接WebSocket
  connectSocket() {
    // 创建WebSocket连接
    wx.connectSocket({
      url: "ws://localhost:8080/ws/ai",
      success: () => {
        console.log("WebSocket连接创建成功")
      },
      fail: (error) => {
        console.error("WebSocket连接创建失败:", error)
        this.handleSocketError()
      },
    })

    // 监听WebSocket连接打开
    wx.onSocketOpen(() => {
      console.log("WebSocket连接已打开")
      this.setData({
        socketConnected: true,
        reconnectCount: 0,
      })
    })

    // 监听WebSocket接收到服务器的消息
    wx.onSocketMessage((res) => {
      this.handleSocketMessage(res)
    })

    // 监听WebSocket错误
    wx.onSocketError((error) => {
      console.error("WebSocket发生错误:", error)
      this.handleSocketError()
    })

    // 监听WebSocket关闭
    wx.onSocketClose((res) => {
      console.log("WebSocket已关闭:", res)
      this.setData({
        socketConnected: false,
      })

      // 如果不是主动停止，尝试重连
      if (!this.data.isStopping && this.data.reconnectCount < this.data.maxReconnectCount) {
        setTimeout(() => {
          this.setData({
            reconnectCount: this.data.reconnectCount + 1,
          })
          this.connectSocket()
        }, 3000) // 3秒后重连
      }
    })
  },

  // 关闭WebSocket连接
  closeSocket() {
    if (this.data.socketConnected) {
      wx.closeSocket({
        success: () => {
          console.log("WebSocket连接已关闭")
        },
      })
    }
  },

  // 处理WebSocket消息
  handleSocketMessage(res) {
    try {
      // 如果正在停止生成，忽略消息
      if (this.data.isStopping) {
        return
      }

      const message = res.data

      // 检查是否包含完成标记
      if (message.includes("[complete]")) {
        // 移除完成标记
        const responseText = this.data.currentResponse + message.replace("[complete]", "")

        // 解析完整的响应
        const parsedResponse = app.towxml(responseText, "markdown", {
          theme: "light",
        })

        // 添加到消息列表
        this.setData(
          {
            messages: [
              ...this.data.messages,
              {
                role: "assistant",
                content: responseText,
                parsedContent: parsedResponse,
                timestamp: new Date().getTime(),
              },
            ],
            loading: false,
            currentResponse: "",
            currentParsedResponse: {},
          },
          () => {
            // 在状态更新完成后滚动到底部
            this.scrollToBottom()
          },
        )
      } else {
        // 累加当前响应
        const responseText = this.data.currentResponse + message

        // 解析当前响应
        const parsedResponse = app.towxml(responseText, "markdown", {
          theme: "light",
        })

        // 更新当前响应
        this.setData(
          {
            currentResponse: responseText,
            currentParsedResponse: parsedResponse,
          },
          () => {
            // 在每次更新后都滚动到底部
            this.scrollToBottom()
          },
        )
      }

      // 滚动到底部
      this.scrollToBottom()
    } catch (error) {
      console.error("处理WebSocket消息失败:", error)
    }
  },

  // 处理WebSocket错误
  handleSocketError() {
    // 如果正在等待回复，显示错误消息
    if (this.data.loading) {
      const errorMessage = "抱歉，连接出现问题，请稍后再试。"
      const parsedErrorMessage = app.towxml(errorMessage, "markdown", {
        theme: "light",
      })

      this.setData({
        messages: [
          ...this.data.messages,
          {
            role: "assistant",
            content: errorMessage,
            parsedContent: parsedErrorMessage,
            timestamp: new Date().getTime(),
            isError: true,
          },
        ],
        loading: false,
        currentResponse: "",
        currentParsedResponse: {},
      })

      this.scrollToBottom()
    }
  },

  // 输入框内容变化
  onInputChange(e) {
    this.setData({
      inputValue: e.detail.value,
    })
  },

  // 发送消息
  sendMessage() {
    const { inputValue, messages, loading, socketConnected } = this.data

    // 如果正在加载或输入为空或WebSocket未连接，不处理
    if (loading || !inputValue.trim() || !socketConnected) {
      if (!socketConnected) {
        wx.showToast({
          title: "网络连接中断，请稍后再试",
          icon: "none",
        })
      }
      return
    }

    // 添加用户消息到列表
    const userMessage = {
      role: "user",
      content: inputValue,
      // 用户消息不需要解析Markdown
      parsedContent: { nodes: [{ type: "text", text: inputValue }] },
      timestamp: new Date().getTime(),
    }

    this.setData({
      messages: [...messages, userMessage],
      inputValue: "",
      loading: true,
      currentResponse: "",
      currentParsedResponse: {},
      isStopping: false,
    })

    // 滚动到底部
    this.scrollToBottom()

    // 通过WebSocket发送消息
    this.sendSocketMessage(inputValue)
  },

  // 停止生成
  stopGeneration() {
    this.setData({
      isStopping: true,
    })

    // 发送停止信号
    this.sendSocketMessage("[stop]")

    // 如果有当前响应，将其添加到消息列表
    if (this.data.currentResponse) {
      const responseText = this.data.currentResponse + " [已停止生成]"
      const parsedResponse = app.towxml(responseText, "markdown", {
        theme: "light",
      })

      this.setData({
        messages: [
          ...this.data.messages,
          {
            role: "assistant",
            content: responseText,
            parsedContent: parsedResponse,
            timestamp: new Date().getTime(),
          },
        ],
        loading: false,
        currentResponse: "",
        currentParsedResponse: {},
      })

      this.scrollToBottom()
    } else {
      this.setData({
        loading: false,
      })
    }
  },

  // 通过WebSocket发送消息
  sendSocketMessage(message) {
    if (!this.data.socketConnected) {
      this.handleSocketError()
      return
    }

    wx.sendSocketMessage({
      data: JSON.stringify({
        prompt: message,
        token: this.data.token || "null"
      }),
      success: () => {
        console.log("消息发送成功")
      },
      fail: (error) => {
        console.error("消息发送失败:", error)
        this.handleSocketError()
      },
    })
  },

  // 滚动到底部
  scrollToBottom() {
    // 使用 nextTick 确保在 DOM 更新后执行滚动
    setTimeout(() => {
      // 获取消息容器
      const query = wx.createSelectorQuery()
      query
        .select("#message-container")
        .fields(
          {
            scrollOffset: true,
            size: true,
          },
          (res) => {
            if (res) {
              // 计算实际需要滚动的位置 - 使用实际内容高度
              const scrollHeight = res.scrollHeight || 0

              this.setData({
                scrollTop: scrollHeight,
              })
            }
          },
        )
        .exec()
    }, 100)
  },

  // 页面显示时滚动到底部并检查连接
  onShow() {
    this.scrollToBottom()

    // 如果WebSocket未连接，尝试重新连接
    if (!this.data.socketConnected) {
      this.connectSocket()
    }
  },

  // 重试发送消息
  retryMessage(e) {
    const { index } = e.currentTarget.dataset
    // 找到对应的用户消息
    let userMessageIndex = index - 1

    // 如果当前是助手消息，需要向前查找最近的用户消息
    if (this.data.messages[index] && this.data.messages[index].role === "assistant") {
      for (let i = index - 1; i >= 0; i--) {
        if (this.data.messages[i] && this.data.messages[i].role === "user") {
          userMessageIndex = i
          break
        }
      }
    }

    const userMessage = this.data.messages[userMessageIndex]

    if (userMessage && userMessage.role === "user") {
      // 移除该用户消息之后的所有消息
      const newMessages = this.data.messages.slice(0, userMessageIndex + 1)

      this.setData({
        messages: newMessages,
        loading: true,
        currentResponse: "",
        currentParsedResponse: {},
        isStopping: false,
      })

      // 重新发送消息
      this.sendSocketMessage(userMessage.content)
    }
  },

  onUnload() {
    // 页面卸载时销毁store实例
    this.tokenStoreBinding.destroyStoreBindings()
  }
})
